package assignment01;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;


import nlp.langmodel.LanguageModel;
import nlp.util.Counter;
import nlp.util.Pair;

/**
 * Start by building a bigram language model using an interpolation method of
 * your choice  Extend you bigram model to a trigram model  Implement either a
 * Katz or Kneser-Ney estimator (bigram ok)
 * 
 * @author chunming
 * 
 */
public class EmpiricalBigramLanguageModel implements LanguageModel {

	private static String START = "<S>";
	private static String STOP = "</S>";
	private static double INCREMENT_FACTOR = 1.0;
//	private Counter<String> wordCounter = new Counter<String>();
	private Counter<Bigram<String, String>> pairCounter = new Counter<Bigram<String, String>>();
	private double pairTotal;
	private Collection<List<String>> validationCollection;
	
	public EmpiricalBigramLanguageModel() {

	}

	public EmpiricalBigramLanguageModel(Collection<List<String>> trainingSentenceCollection, Collection<List<String>> validationSentenceCollection) {
		this.validationCollection = validationSentenceCollection;
		for (List<String> sentence : trainingSentenceCollection) {
			List<String> stopSentence = new ArrayList<String>(sentence);
			stopSentence.add(0, START);
			stopSentence.add(STOP);
//			for (String word : stopSentence) {
//				wordCounter.incrementCount(word, INCREMENT_FACTOR);
//			}
			Iterator<String> iter = stopSentence.iterator();
			String first = iter.next();
			String second = iter.next();

			while (iter.hasNext() || !second.equals(STOP)) {
				Bigram bigram = new Bigram(first, second);
				pairCounter.incrementCount(bigram, INCREMENT_FACTOR);
				first = second;
				if (iter.hasNext()) {
					second = iter.next();
				}
			}
		}
//		wordTotal = wordCounter.totalCount();
		pairTotal = pairCounter.totalCount();
//		System.out.println("wordTotal is " + wordTotal + ", and pairTotal is " + pairTotal);
//		System.out.println("(a, share) occurance is " + pairCounter.getCount(new Bigram("a", "share")));
//		System.out.println("(share, of) occurance is " + pairCounter.getCount(new Bigram("share", "of")));
//		System.out.println("(of, the) occurance is " + pairCounter.getCount(new Bigram("of", "the")));
//		System.out.println("(the, company) occurance is " + pairCounter.getCount(new Bigram("the", "company")));
//		System.out.println("(compay, 's) occurance is " + pairCounter.getCount(new Bigram("company", "'s")));
//		System.out.println("('s, a) occurance is " + pairCounter.getCount(new Bigram("'s", "a")));
//		System.out.println("pairCounter.keySet().size() " + pairCounter.keySet().size());
		
	}

	private String generateStartWord() {
		double threshold = Math.random();
		double sum = 0.0;
		for (Bigram bigram: pairCounter.keySet()) {
			sum += pairCounter.getCount(bigram) / pairTotal;
			if (bigram.getFirst().equals(START) && sum >= threshold) {
				System.out.println("second is" + bigram.getSecond().toString());
				return (String)bigram.getSecond();
			}
		}
//		return null;
		return STOP;
	}
	
	
	/**
	 * Maximum Likely Estimation
	 * @param previous
	 * @return
	 */
	private String generateWord(String previous) {
		System.out.println("previously is " + previous);
		double likelihood = 0.0;
		Bigram likeliBigram = null;
		for (Bigram bigram: pairCounter.keySet()) {
//			System.out.println(bigram);
			if (bigram.getFirst().toString().equalsIgnoreCase(previous)) {
				if (likelihood < pairCounter.getCount(bigram)) {
					likelihood = pairCounter.getCount(bigram);
					likeliBigram = bigram;
				}
				if (likeliBigram.getSecond().equals(STOP)) {
					return likeliBigram.getSecond().toString();
				}
			}
		}
		if (likeliBigram != null) {
			return likeliBigram.getSecond().toString();
		}
		return null;
	}
	
	private double getWordProbability(List<String> sentence, int index) {
		if (index <= 0) {
			return 1.0;
		} else {
			String word = sentence.get(index);
			String previous = sentence.get(index - 1);
			Bigram<String, String> bigram = new Bigram<String, String>(previous, word);
			double count = pairCounter.getCount(bigram);
			double probability = count == 0.0 ? 1.0 : count / pairTotal;
//			System.out.println("word probability is " + probability);

			return probability;
		}
	}
	
	@Override
	public double getSentenceProbability(List<String> sentence) {
		List<String> stoppedSentence = new ArrayList<String>(sentence);
		stoppedSentence.add(0, START);
		stoppedSentence.add(STOP);
		double probability = 0.0;
		for (int i = 0; i < stoppedSentence.size(); i++) {
//			System.out.println("the log word probability is " + Math.log(getWordProbability(stoppedSentence, i)));
			probability += Math.log(getWordProbability(stoppedSentence, i));
		}
//		probability = Math.pow(Math.E, probability);
		return probability;
	}
	
	public double[] solveLambda() {
		double[] container = new double[100];
		for (int i = 0; i < container.length; i++) {
			container[i] = LanguageModelTester.calculatePerplexity(this, validationCollection);
			System.out.println("perplexity i: " + container[i]);
		}
		return container;
	}

	@Override
	public List<String> generateSentence() {
		solveLambda();
		List<String> sentence = new ArrayList<String>();
		sentence.add(START);
		String word = generateWord(START);
//		String word = generateStartWord();
		while (!word.equals(STOP)) {
			sentence.add(word);
			word = generateWord(word);
		}
		return sentence;
	}
}
